package org.xx.armory.commons;

import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.openssl.PEMEncryptedKeyPair;
import org.bouncycastle.openssl.PEMKeyPair;
import org.bouncycastle.openssl.PEMParser;
import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;
import org.bouncycastle.util.io.pem.PemObject;
import org.bouncycastle.util.io.pem.PemWriter;
import org.slf4j.LoggerFactory;

import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.DESKeySpec;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.UncheckedIOException;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

import static org.xx.armory.commons.Validators.rejectIfBlank;
import static org.xx.armory.commons.Validators.rejectIfNull;

/**
 * 密钥相关的工具程序。
 * <p>允许从PEM格式的文件中读取私钥、公钥或者证书。</p>
 *
 * @see <a href="https://tools.ietf.org/html/rfc1421" >RFC1421</a>
 */
public final class KeyUtils {
    static {
        final var logger = LoggerFactory.getLogger(KeyUtils.class);

        final var _BCProviderFullClassName = "org.bouncycastle.jce.provider.BouncyCastleProvider";
        try {
            final var bouncyCastleProviderClass = Class.forName(_BCProviderFullClassName).asSubclass(java.security.Provider.class);
            java.security.Security.addProvider(bouncyCastleProviderClass.getConstructor().newInstance());
            logger.info("Security provider {} loaded", _BCProviderFullClassName);
        } catch (ReflectiveOperationException ex) {
            logger.error("Cannot load: " + _BCProviderFullClassName, ex);
        }
    }

    private KeyUtils() {
        throw new AssertionError();
    }

    /**
     * 从PEM格式的资源文件中读取私钥。
     *
     * @param stream
     *         包含PEM格式的私钥的输入流。
     * @return 已读取的私钥。
     * @throws IllegalArgumentException
     *         如果参数{@code resource}是{@code null}或者不可读。
     * @throws UncheckedIOException
     *         如果无法打开资源文件并读取内容，如果资源文件的内容不是合法的PEM格式，如果PEM文件中的私钥被加密。
     */
    public static PrivateKey loadPrivateKeyFromPEM(
            InputStream stream
    ) {
        rejectIfNull(stream, "stream");

        try (final var reader = new InputStreamReader(stream);
             final var pemParser = new PEMParser(reader)) {
            final var object = pemParser.readObject();
            if (object instanceof PEMKeyPair) {
                final var kp = new JcaPEMKeyConverter().setProvider("BC").getKeyPair((PEMKeyPair) object);
                return kp.getPrivate();
            } else if (object instanceof PrivateKey) {
                return (PrivateKey) object;
            } else if (object instanceof PrivateKeyInfo) {
                return new JcaPEMKeyConverter().setProvider("BC").getPrivateKey((PrivateKeyInfo) object);
            } else if (object instanceof PEMEncryptedKeyPair) {
                throw new IOException("Cannot parse encrypted key");
            } else {
                throw new IOException("Cannot find key-pair, got " + object);
            }
        } catch (IOException ex) {
            throw new UncheckedIOException(ex);
        }
    }

    /**
     * 从PEM格式的资源文件中读取公钥。
     *
     * @param stream
     *         包含PEM格式的公钥的输入流。
     * @return 已读取的公钥。
     * @throws IllegalArgumentException
     *         如果参数{@code resource}是{@code null}或者不可读。
     * @throws UncheckedIOException
     *         如果无法打开资源文件并读取内容，如果资源文件的内容不是合法的PEM格式。
     */
    public static PublicKey loadPublicKeyFromPEM(
            InputStream stream
    ) {
        rejectIfNull(stream, "stream");

        try (final var reader = new InputStreamReader(stream);
             final var pemParser = new PEMParser(reader)) {
            final var object = pemParser.readObject();
            if (object instanceof PEMKeyPair) {
                final var kp = new JcaPEMKeyConverter().setProvider("BC").getKeyPair((PEMKeyPair) object);
                return kp.getPublic();
            } else if (object instanceof PublicKey) {
                return (PublicKey) object;
            } else if (object instanceof SubjectPublicKeyInfo) {
                return new JcaPEMKeyConverter().setProvider("BC").getPublicKey((SubjectPublicKeyInfo) object);
            } else if (object instanceof PEMEncryptedKeyPair) {
                throw new IOException("Cannot parse encrypted key");
            } else {
                throw new IOException("Cannot find key-pair, got " + object);
            }
        } catch (IOException ex) {
            throw new UncheckedIOException(ex);
        }
    }

    /**
     * 从PEM格式的资源文件中读取证书。
     *
     * @param stream
     *         包含PEM格式的证书的输入流。
     * @return 已读取的证书。
     * @throws IllegalArgumentException
     *         如果参数{@code resource}是{@code null}或者不可读。
     * @throws IllegalStateException
     *         如果当前jvm无法解析x.509格式的证书。
     */
    public static Certificate loadCertificateFromPEM(
            InputStream stream
    ) {
        rejectIfNull(stream, "stream");

        try {
            return CertificateFactory.getInstance("X.509").generateCertificate(stream);
        } catch (CertificateException ex) {
            throw new IllegalStateException("Cannot parse x.509 pem", ex);
        }
    }

    private static KeyFactory createRSAKeyFactory() {
        try {
            return KeyFactory.getInstance("RSA");
        } catch (NoSuchAlgorithmException ex) {
            throw new RuntimeException("Cannot get RSA algorithm", ex);
        }
    }

    private static SecretKeyFactory createDESKeyFactory() {
        try {
            return SecretKeyFactory.getInstance("DES");
        } catch (NoSuchAlgorithmException ex) {
            throw new RuntimeException("Cannot get DES algorithm", ex);
        }
    }

    /**
     * 将DER格式的公钥转化为PEM格式。
     *
     * @param s
     *         Base64格式的公钥。
     * @return PEM格式的公钥证书。
     * @throws IllegalStateException
     *         如果当前jvm无法解析x.509格式的证书。
     */
    public static PublicKey loadRSAPublicKeyFromX509DER(
            String s
    ) {
        s = rejectIfBlank(s, "s").strip();

        try {
            final var bytes = Base64.getDecoder().decode(s);
            return createRSAKeyFactory().generatePublic(new X509EncodedKeySpec(bytes));
        } catch (InvalidKeySpecException ex) {
            throw new IllegalStateException("Cannot parse x.509 der", ex);
        }
    }

    /**
     * 加载DER格式的私钥。
     *
     * @param s
     *         Base64格式的私钥。
     * @return PEM格式的私钥证书。
     * @throws IllegalStateException
     *         如果当前jvm无法解析PKCS8格式的证书。
     */
    public static PrivateKey loadRSAPrivateKeyFromPKCS8DER(
            String s
    ) {
        s = rejectIfBlank(s, "s").strip();

        try {
            final var bytes = Base64.getDecoder().decode(s);
            return createRSAKeyFactory().generatePrivate(new PKCS8EncodedKeySpec(bytes));
        } catch (InvalidKeySpecException ex) {
            throw new IllegalStateException("Cannot parse PKCS8 der", ex);
        }
    }

    /**
     * 加载Base64编码的密钥。
     *
     * @param s
     *         Base64格式的密钥。
     * @return 密钥。
     * @throws IllegalStateException
     *         如果当前jvm无法解析Base64格式的密钥。
     */
    public static SecretKey loadDESKeyFromBase64(
            String s
    ) {
        s = rejectIfBlank(s, "s").strip();

        try {
            final var bytes = Base64.getDecoder().decode(s);
            return createDESKeyFactory().generateSecret(new DESKeySpec(bytes));
        } catch (InvalidKeySpecException | InvalidKeyException ex) {
            throw new IllegalStateException("Cannot parse DES", ex);
        }
    }

    public static void main(String[] args)
            throws Exception {
        final var s = "MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA21ekeivwf8hP6Ic7RNT9UsJPzS5yb83kQHZKne2+GrWjaDwQIgQpguGsjb1LfLiAeTz4GrhWfayXhmisz2McaffRlKlNI/7GTi6PpdP9nTjk8Ihj3vakePgPLCMksqEDTD9rYxg0ShyY0/My6P8tCcmw83AP2o+2YI9JsoaoY/YixpY6aEkCQomJZugja7nAnLPEcIqMWyCIOjzdz+NfI5nElEALDpf3oE3vDO2WRbiPxO45X/KOHG4fRvQjork1FRrbs1adIgLU0+yt+y3DRJyfVD2vyuISNEq/OvOgv3Xo8+SQaWbPG0/DmmWS6TQ+8sjtShUUfqbhp1ZOq8JD8wIDAQAB";

        try (final var pw = new PemWriter(new FileWriter("a.pem"))) {
            pw.writeObject(new PemObject("PUBLIC KEY", loadRSAPublicKeyFromX509DER(s).getEncoded()));
        }

        final var p = "MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDbV6R6K/B/yE/ohztE1P1Swk/NLnJvzeRAdkqd7b4ataNoPBAiBCmC4ayNvUt8uIB5PPgauFZ9rJeGaKzPYxxp99GUqU0j/sZOLo+l0/2dOOTwiGPe9qR4+A8sIySyoQNMP2tjGDRKHJjT8zLo/y0JybDzcA/aj7Zgj0myhqhj9iLGljpoSQJCiYlm6CNrucCcs8RwioxbIIg6PN3P418jmcSUQAsOl/egTe8M7ZZFuI/E7jlf8o4cbh9G9COiuTUVGtuzVp0iAtTT7K37LcNEnJ9UPa/K4hI0Sr8686C/dejz5JBpZs8bT8OaZZLpND7yyO1KFRR+puGnVk6rwkPzAgMBAAECggEBAK0D6fKQnEk/SYDiwuadlyDjsyzm2U/Fuaj7ZgpDEPRimH3M+vDKFbBusLyxN00u//AiHpZ9F4RUUsWr5bxEspIH7+Jy1S1ydxwOcDLfECsnSLdO7nBeXwgoq1cmdXkNtts7H0rwObbH7ekp6duOBQm2dYTgvTCqsbJi9EAFJeiKNZMGUzZLOcdfu5yWeqs1hF/poNNY/3TA/ueQnxE9MLvQBMGvf+tA1M3lniiUGLmONSfrKKMLC7gbSEPYQCix/S+k4a3QiWyyIx3EqkIRuvs6Sf/rqeq1POD/yIumpbuEgSSUBHw/pkqOnHEroZwHu562QsYbxGHMFkjWap7kE8ECgYEA867vkm6oPiwHp2J2hohnpFaqVAUKozoP3COtOt5xwpwFYaGTo/Nix5A5NUhVi7DSumNOVCaZnkeMJSQjny6bPVKCRZ+nalKYquuwTUmA3UYiBW4fWZJrjk7u9z42CBaYkjkT79csr0mlImr9jIO9sdyrXXV2ouGzzLDK5ybVyJMCgYEA5m3BCjZWvwimFoIeSsdC1xSUey+axuLu/VRzOl1p5Wt+Dz0czFBlEc0+uW4jc7ZD+Y0QJHYtCFGuYQuhjm/04CzMhwdlpz+Sku6o6xAIf2Gee64U+mQ5XlKwLreHXmHZmL7/qsk+lV8ZgfaPgSBhGGW2J3Z04hZ9+8v3uY89kyECgYEA3sOyALQmG0fW616oqY2y6VnMtagyNEtUylLW/rTxXA2TIdA6jqtlRrgChEIO6feIfJlhtq3eywHvDNj/s3pAsgtmSc2MlzbQl5eEAPNL2XxOh5ZHFHQ03NAtxlnZvQUjDrEdF6kmEUwKqO8TwEuCs0Deaz6BaJ9Z5apMooku7/8CgYAksIcTHDHVuwfO/6rZFz8iRemab6sCTfHlrkOmoGcTG3l37bAfLBHbssVw+wN0Eipimr5eP97Nu0u7sgFoT5U6+IccXOYnqiypbozA2M2xkpaMdNxsPfcQeLQKzDgmpAfklUzaIa7wU7MZCWwyEDjFjLRjNi5ILXKj3Ac2qYMRYQKBgEJ+H8UYzKuKBQ3sPe9BD9sEsJVGrD3/Vfq8XAOmSFT5WA0/uwBl45Z92gYwz5YfRRoq0LcSEVhKWSyxYHMXuSk0PsjmU7QyUJsqIwlLcQtjA4RWsUU3PvBQMLZE+C8NEEpT0m/MqkyPqRvF0C7RrZeLc70Y0vEICfUI8C804veS";
        try (final var pw = new PemWriter(new FileWriter("a.key"))) {
            pw.writeObject(new PemObject("PRIVATE KEY", loadRSAPrivateKeyFromPKCS8DER(p).getEncoded()));
        }
    }
}
